import os
import sys
import time
import urllib.parse
import shutil
import requests
import re
from pathlib import Path

# 配置区
ROOT_DIR = "/workspace"
DOWNLOAD_DIR = os.path.join(ROOT_DIR, "URLDown")
MAX_RETRY = 3
TIMEOUT = 30

# 颜色定义
RED = "\033[31m"
GREEN = "\033[32m"
YELLOW = "\033[33m"
BLUE = "\033[34m"
CYAN = "\033[36m"
PURPLE = "\033[35m"
RESET = "\033[0m"

# 创建下载目录
os.makedirs(DOWNLOAD_DIR, exist_ok=True)

from urllib.parse import urlparse, urlunparse

def transform_huggingface_url(url: str) -> str:
    if url.startswith("https://huggingface.co/") or url.startswith("https://hf-mirror.com/"):
        parsed = urlparse(url)
        parts = parsed.path.strip("/").split("/")
        if len(parts) > 2:
            parts[2] = "resolve"

        # 替换 netloc 并重组路径
        new_path = "/" + "/".join(parts)
        new_parsed = parsed._replace(netloc="hf-mirror.com", path=new_path)
        return urlunparse(new_parsed)
    return url


def extract_filename_from_url(url: str) -> str:
    parsed_url = urllib.parse.urlparse(url)
    path = parsed_url.path
    basename = os.path.basename(path)
    filename = urllib.parse.unquote(basename)

    if '.' in filename and not filename.startswith('.'):
        return filename
    else:
        return f"{filename}.download"

def check_tool_available(tool: str) -> bool:
    """检查下载工具是否可用"""
    try:
        if tool == "curl":
            result = os.system("curl --version >nul 2>&1")
            return result == 0
        elif tool == "wget":
            result = os.system("wget --version >nul 2>&1")
            return result == 0
        else:
            return True  # requests总是可用的
    except:
        return False

def select_download_tool() -> str:
    print(f"\n{CYAN}请选择下载工具：{RESET}")
    
    # 检查工具可用性
    curl_available = check_tool_available("curl")
    wget_available = check_tool_available("wget")
    
    print("1) requests（推荐，总是可用）")
    if curl_available:
        print("2) curl（系统支持）")
    else:
        print("2) curl（不可用）")
    if wget_available:
        print("3) wget（系统支持）")
    else:
        print("3) wget（不可用，需要安装）")
    
    choice = input("输入选项 (1/2/3): ").strip()
    
    if choice == "2" and curl_available:
        return "curl"
    elif choice == "3" and wget_available:
        return "wget"
    elif choice == "2" and not curl_available:
        print(f"{RED}curl不可用，使用requests{RESET}")
        return "requests"
    elif choice == "3" and not wget_available:
        print(f"{RED}wget不可用，使用requests{RESET}")
        return "requests"
    else:
        print(f"{YELLOW}默认使用 requests{RESET}")
        return "requests"

def download_file(url: str, filename: str, tool: str):
    url = transform_huggingface_url(url)
    auto_filename = extract_filename_from_url(url)

    if not filename:
        filename = auto_filename
        print(f"{YELLOW}未指定文件名，自动使用: {filename}{RESET}")
    elif '.' not in filename:
        ext = os.path.splitext(auto_filename)[1]
        filename += ext
        print(f"{YELLOW}已添加扩展名: {filename}{RESET}")

    filepath = os.path.join(DOWNLOAD_DIR, filename)

    print(f"\n{BLUE}■ 下载信息 ■{RESET}")
    print(f"URL: {PURPLE}{url}{RESET}")
    print(f"保存位置: {PURPLE}{filepath}{RESET}")

    print(f"\n{CYAN}开始下载...{RESET}")
    success = False

    for attempt in range(1, MAX_RETRY + 1):
        print(f"{BLUE}▶ 尝试第 {attempt} 次 (最多 {MAX_RETRY} 次){RESET}")
        try:
            if tool == "curl":
                print(f"{CYAN}使用curl下载...{RESET}")
                result = os.system(f'curl -L -C - --connect-timeout {TIMEOUT} --progress-bar -o "{filepath}" "{url}"')
                success = result == 0 and os.path.exists(filepath)
            elif tool == "wget":
                print(f"{CYAN}使用wget下载...{RESET}")
                result = os.system(f'wget -c --timeout={TIMEOUT} --progress=bar -O "{filepath}" "{url}"')
                success = result == 0 and os.path.exists(filepath)
            else:
                # 改进的requests下载
                headers = {
                    'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
                }
                
                with requests.get(url, stream=True, timeout=TIMEOUT, headers=headers) as r:
                    r.raise_for_status()
                    total_size = int(r.headers.get('content-length', 0))
                    
                    print(f"{CYAN}文件大小: {total_size / (1024*1024):.2f} MB{RESET}")
                    
                    with open(filepath, 'wb') as f:
                        downloaded = 0
                        for chunk in r.iter_content(chunk_size=8192):
                            if chunk:
                                f.write(chunk)
                                downloaded += len(chunk)
                                if total_size > 0:
                                    percent = (downloaded / total_size) * 100
                                    print(f"\r{CYAN}下载进度: {percent:.1f}% ({downloaded/(1024*1024):.2f}MB/{total_size/(1024*1024):.2f}MB){RESET}", end='', flush=True)
                    
                    print()  # 换行
                success = True

            if success:
                print(f"\n{CYAN}当前磁盘使用情况：{RESET}")
                os.system("df -h")

                break
        except Exception as e:
            print(f"{YELLOW}⚠ 下载失败：{e}，等待重试...{RESET}")
            time.sleep(3)

    if success:
        print(f"\n{GREEN}✓ 下载成功！{RESET}")
        try:
            size = os.path.getsize(filepath)
            print(f"{BLUE}文件大小: {size / 1024:.2f} KB{RESET}")
        except:
            pass
    else:
        print(f"\n{RED}✗ 下载失败！{RESET}")
        if os.path.exists(filepath):
            os.remove(filepath)

def main():
    print(f"\n{PURPLE}■■■ URL下载器 v3.0（Python版） ■■■{RESET}")
    print(f"下载目录: {CYAN}{os.path.abspath(DOWNLOAD_DIR)}{RESET}")
    
    tool = select_download_tool()

    while True:
        print(f"\n{BLUE}════════════════════════════{RESET}")
        url = input("请输入下载URL (或输入 'q' 退出): ").strip()
        if url.lower() == "q":
            print(f"{GREEN}感谢使用，再见！{RESET}")
            break
        elif not url:
            print(f"{RED}错误：URL不能为空！{RESET}")
            continue

        filename = input("请输入保存文件名 (留空自动从URL提取): ").strip()
        download_file(url, filename, tool)

if __name__ == "__main__":
    main()
